import logging

import torch
import sys


sys.path.insert(0,'../../')
from gxl_ai_utils.utils import utils_file
from types import SimpleNamespace
from wenet.zipformer.init_zipformer import init_zipformer2_encoder
from wenet.utils.checkpoint import load_checkpoint
from wenet.utils.init_model import init_model
from wenet.zipformer.init_decoder_attention import init_decoder4zipformer
from wenet.zipformer.config import CONFIG_ZIPFORMER

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(levelname)s %(message)s')
def do_test_encoder():
    model_pt_path = "/home/work_nfs15/asr_data/ckpt/asr_online_system/epoch-10.pt"
    print(model_pt_path)
    checkpoint = torch.load(model_pt_path, map_location='cpu')
    params = checkpoint['model']
    new_param_dict_only_encoder= {}
    for k, v in params.items():
        if k.startswith("encoder"):
            new_param_dict_only_encoder[k] = v
    # for k in new_param_dict_only_encoder.keys():
    #     print(k)
    configs = utils_file.load_dict_from_yaml('train_zipformer.yaml')
    args = SimpleNamespace(**{
        'name':'耿雪龙'
    })
    print(args.name)
    model, configs = init_model(args, configs)
    missing_keys, unexpected_keys = model.load_state_dict(new_param_dict_only_encoder,
                                                          strict=False)
    for key in missing_keys:
        logging.info("missing tensor: {}".format(key))
    for key in unexpected_keys:
        logging.info("unexpected tensor: {}".format(key))


def do_test_encoder_embeded():
    encoder = init_zipformer2_encoder()
    print(encoder)
    input_wav = torch.randn(2, 32, 80)
    input_wav_lens = torch.tensor([32,21]) # ((time-4) /// 2 - 2)///2    ; ///代表除以后取上界
    encoder_out, encoder_out_lens = encoder(input_wav, input_wav_lens)
    print(encoder_out.shape)
    print(encoder_out_lens)
    model_pt_path = "/home/work_nfs15/asr_data/ckpt/asr_online_system/epoch-10.pt"
    print(model_pt_path)
    checkpoint = torch.load(model_pt_path, map_location='cpu')
    params = checkpoint['model']
    new_param_dict_only_encoder = {}
    for k, v in params.items():
        if k.startswith("encoder."):
            new_param_dict_only_encoder[k] = v
        elif k.startswith("encoder_embed."):
            k = k.replace("encoder_embed.", "encoder.encoder_embed.")
            new_param_dict_only_encoder[k] = v
    configs = utils_file.load_dict_from_yaml('train_zipformer.yaml')
    args = SimpleNamespace(**{
        'name': '耿雪龙'
    })
    print(args.name)
    model, configs = init_model(args, configs)
    missing_keys, unexpected_keys = model.load_state_dict(new_param_dict_only_encoder,
                                                          strict=False)
    for key in missing_keys:
        logging.info("missing tensor: {}".format(key))
    for key in unexpected_keys:
        logging.info("unexpected tensor: {}".format(key))

def do_test_encoder_size_only():
    encoder = init_zipformer2_encoder(CONFIG_ZIPFORMER)
    utils_file.print_model_size(encoder)

def do_test_decoder_size_only():
    """"""
    decoder = init_decoder4zipformer(CONFIG_ZIPFORMER)
    utils_file.print_model_size(decoder)

def do_test_full_model():
    """"""
    model_pt_path = "/home/work_nfs15/asr_data/ckpt/asr_online_system/epoch-10.pt"
    print(model_pt_path)
    checkpoint = torch.load(model_pt_path, map_location='cpu')
    params = checkpoint['model']
    new_param_dict_only_encoder_decoder = {}
    for k, v in params.items():
        if k.startswith("encoder."):
            new_param_dict_only_encoder_decoder[k] = v
        elif k.startswith("attention_decoder."):
            k = k.replace("attention_decoder.", "")
            new_param_dict_only_encoder_decoder[k] = v
        elif k.startswith("encoder_embed."):
            k = k.replace("encoder_embed.", "encoder.encoder_embed.")
            new_param_dict_only_encoder_decoder[k] = v
        elif k.startswith("ctc_output.1"):
            k = k.replace("ctc_output.1", "ctc.ctc_lo")
            new_param_dict_only_encoder_decoder[k] = v
    output_model_pt_path = "/home/work_nfs15/asr_data/ckpt/asr_online_system/epoch-10_model.pt"
    torch.save(params, output_model_pt_path)
    configs = utils_file.load_dict_from_yaml('train_zipformer.yaml')
    args = SimpleNamespace(**{
        'name': '耿雪龙'
    })
    print(args.name)
    model, configs = init_model(args, configs)
    missing_keys, unexpected_keys = model.load_state_dict(params,
                                                          strict=False)
    for key in missing_keys:
        logging.info("missing tensor: {}".format(key))
    for key in unexpected_keys:
        logging.info("unexpected tensor: {}".format(key))

if __name__ == '__main__':
    do_test_full_model()

